import mindspore as ms

convert_dict = {
    "feature_map.backbone.conv0.1": "model.model.0.bn",
    "feature_map.backbone.conv1.1": "model.model.1.bn",
    "feature_map.backbone.conv2.1": "model.model.2.bn",
    "feature_map.backbone.conv3.1": "model.model.3.conv1.bn",
    "feature_map.backbone.conv4.1": "model.model.3.conv2.bn",
    "feature_map.backbone.conv5.1": "model.model.4.bn",
    "feature_map.backbone.conv6.1": "model.model.5.bn",
    "feature_map.backbone.conv7.1": "model.model.7.bn",
    "feature_map.backbone.conv8.1": "model.model.8.bn",
    "feature_map.backbone.conv9.1": "model.model.9.bn",
    "feature_map.backbone.layer2.0.conv1.1": "model.model.10.0.conv1.bn",
    "feature_map.backbone.layer2.0.conv2.1": "model.model.10.0.conv2.bn",
    "feature_map.backbone.layer2.1.conv1.1": "model.model.10.1.conv1.bn",
    "feature_map.backbone.layer2.1.conv2.1": "model.model.10.1.conv2.bn",
    "feature_map.backbone.conv10.1": "model.model.11.bn",
    "feature_map.backbone.conv11.1": "model.model.12.bn",
    "feature_map.backbone.conv12.1": "model.model.14.bn",
    "feature_map.backbone.conv13.1": "model.model.15.bn",
    "feature_map.backbone.conv14.1": "model.model.16.bn",
    "feature_map.backbone.layer3.0.conv1.1": "model.model.17.0.conv1.bn",
    "feature_map.backbone.layer3.1.conv1.1": "model.model.17.1.conv1.bn",
    "feature_map.backbone.layer3.2.conv1.1": "model.model.17.2.conv1.bn",
    "feature_map.backbone.layer3.3.conv1.1": "model.model.17.3.conv1.bn",
    "feature_map.backbone.layer3.4.conv1.1": "model.model.17.4.conv1.bn",
    "feature_map.backbone.layer3.5.conv1.1": "model.model.17.5.conv1.bn",
    "feature_map.backbone.layer3.6.conv1.1": "model.model.17.6.conv1.bn",
    "feature_map.backbone.layer3.7.conv1.1": "model.model.17.7.conv1.bn",
    "feature_map.backbone.layer3.0.conv2.1": "model.model.17.0.conv2.bn",
    "feature_map.backbone.layer3.1.conv2.1": "model.model.17.1.conv2.bn",
    "feature_map.backbone.layer3.2.conv2.1": "model.model.17.2.conv2.bn",
    "feature_map.backbone.layer3.3.conv2.1": "model.model.17.3.conv2.bn",
    "feature_map.backbone.layer3.4.conv2.1": "model.model.17.4.conv2.bn",
    "feature_map.backbone.layer3.5.conv2.1": "model.model.17.5.conv2.bn",
    "feature_map.backbone.layer3.6.conv2.1": "model.model.17.6.conv2.bn",
    "feature_map.backbone.layer3.7.conv2.1": "model.model.17.7.conv2.bn",
    "feature_map.backbone.conv15.1": "model.model.18.bn",
    "feature_map.backbone.conv16.1": "model.model.19.bn",
    "feature_map.backbone.conv17.1": "model.model.21.bn",
    "feature_map.backbone.conv18.1": "model.model.22.bn",
    "feature_map.backbone.conv19.1": "model.model.23.bn",
    "feature_map.backbone.layer4.0.conv1.1": "model.model.24.0.conv1.bn",
    "feature_map.backbone.layer4.1.conv1.1": "model.model.24.1.conv1.bn",
    "feature_map.backbone.layer4.2.conv1.1": "model.model.24.2.conv1.bn",
    "feature_map.backbone.layer4.3.conv1.1": "model.model.24.3.conv1.bn",
    "feature_map.backbone.layer4.4.conv1.1": "model.model.24.4.conv1.bn",
    "feature_map.backbone.layer4.5.conv1.1": "model.model.24.5.conv1.bn",
    "feature_map.backbone.layer4.6.conv1.1": "model.model.24.6.conv1.bn",
    "feature_map.backbone.layer4.7.conv1.1": "model.model.24.7.conv1.bn",
    "feature_map.backbone.layer4.0.conv2.1": "model.model.24.0.conv2.bn",
    "feature_map.backbone.layer4.1.conv2.1": "model.model.24.1.conv2.bn",
    "feature_map.backbone.layer4.2.conv2.1": "model.model.24.2.conv2.bn",
    "feature_map.backbone.layer4.3.conv2.1": "model.model.24.3.conv2.bn",
    "feature_map.backbone.layer4.4.conv2.1": "model.model.24.4.conv2.bn",
    "feature_map.backbone.layer4.5.conv2.1": "model.model.24.5.conv2.bn",
    "feature_map.backbone.layer4.6.conv2.1": "model.model.24.6.conv2.bn",
    "feature_map.backbone.layer4.7.conv2.1": "model.model.24.7.conv2.bn",
    "feature_map.backbone.conv20.1": "model.model.25.bn",
    "feature_map.backbone.conv21.1": "model.model.26.bn",
    "feature_map.backbone.conv22.1": "model.model.28.bn",
    "feature_map.backbone.conv23.1": "model.model.29.bn",
    "feature_map.backbone.conv24.1": "model.model.30.bn",
    "feature_map.backbone.layer5.0.conv1.1": "model.model.31.0.conv1.bn",
    "feature_map.backbone.layer5.1.conv1.1": "model.model.31.1.conv1.bn",
    "feature_map.backbone.layer5.2.conv1.1": "model.model.31.2.conv1.bn",
    "feature_map.backbone.layer5.3.conv1.1": "model.model.31.3.conv1.bn",
    "feature_map.backbone.layer5.0.conv2.1": "model.model.31.0.conv2.bn",
    "feature_map.backbone.layer5.1.conv2.1": "model.model.31.1.conv2.bn",
    "feature_map.backbone.layer5.2.conv2.1": "model.model.31.2.conv2.bn",
    "feature_map.backbone.layer5.3.conv2.1": "model.model.31.3.conv2.bn",
    "feature_map.backbone.conv25.1": "model.model.32.bn",
    "feature_map.backbone.conv26.1": "model.model.33.bn",
    "feature_map.backbone.conv27.1": "model.model.35.bn",
    "feature_map.backbone.conv0.0": "model.model.0.conv",
    "feature_map.backbone.conv1.0": "model.model.1.conv",
    "feature_map.backbone.conv2.0": "model.model.2.conv",
    "feature_map.backbone.conv3.0": "model.model.3.conv1.conv",
    "feature_map.backbone.conv4.0": "model.model.3.conv2.conv",
    "feature_map.backbone.conv5.0": "model.model.4.conv",
    "feature_map.backbone.conv6.0": "model.model.5.conv",
    "feature_map.backbone.conv7.0": "model.model.7.conv",
    "feature_map.backbone.conv8.0": "model.model.8.conv",
    "feature_map.backbone.conv9.0": "model.model.9.conv",
    "feature_map.backbone.layer2.0.conv1.0": "model.model.10.0.conv1.conv",
    "feature_map.backbone.layer2.0.conv2.0": "model.model.10.0.conv2.conv",
    "feature_map.backbone.layer2.1.conv1.0": "model.model.10.1.conv1.conv",
    "feature_map.backbone.layer2.1.conv2.0": "model.model.10.1.conv2.conv",
    "feature_map.backbone.conv10.0": "model.model.11.conv",
    "feature_map.backbone.conv11.0": "model.model.12.conv",
    "feature_map.backbone.conv12.0": "model.model.14.conv",
    "feature_map.backbone.conv13.0": "model.model.15.conv",
    "feature_map.backbone.conv14.0": "model.model.16.conv",
    "feature_map.backbone.layer3.0.conv1.0": "model.model.17.0.conv1.conv",
    "feature_map.backbone.layer3.1.conv1.0": "model.model.17.1.conv1.conv",
    "feature_map.backbone.layer3.2.conv1.0": "model.model.17.2.conv1.conv",
    "feature_map.backbone.layer3.3.conv1.0": "model.model.17.3.conv1.conv",
    "feature_map.backbone.layer3.4.conv1.0": "model.model.17.4.conv1.conv",
    "feature_map.backbone.layer3.5.conv1.0": "model.model.17.5.conv1.conv",
    "feature_map.backbone.layer3.6.conv1.0": "model.model.17.6.conv1.conv",
    "feature_map.backbone.layer3.7.conv1.0": "model.model.17.7.conv1.conv",
    "feature_map.backbone.layer3.0.conv2.0": "model.model.17.0.conv2.conv",
    "feature_map.backbone.layer3.1.conv2.0": "model.model.17.1.conv2.conv",
    "feature_map.backbone.layer3.2.conv2.0": "model.model.17.2.conv2.conv",
    "feature_map.backbone.layer3.3.conv2.0": "model.model.17.3.conv2.conv",
    "feature_map.backbone.layer3.4.conv2.0": "model.model.17.4.conv2.conv",
    "feature_map.backbone.layer3.5.conv2.0": "model.model.17.5.conv2.conv",
    "feature_map.backbone.layer3.6.conv2.0": "model.model.17.6.conv2.conv",
    "feature_map.backbone.layer3.7.conv2.0": "model.model.17.7.conv2.conv",
    "feature_map.backbone.conv15.0": "model.model.18.conv",
    "feature_map.backbone.conv16.0": "model.model.19.conv",
    "feature_map.backbone.conv17.0": "model.model.21.conv",
    "feature_map.backbone.conv18.0": "model.model.22.conv",
    "feature_map.backbone.conv19.0": "model.model.23.conv",
    "feature_map.backbone.layer4.0.conv1.0": "model.model.24.0.conv1.conv",
    "feature_map.backbone.layer4.1.conv1.0": "model.model.24.1.conv1.conv",
    "feature_map.backbone.layer4.2.conv1.0": "model.model.24.2.conv1.conv",
    "feature_map.backbone.layer4.3.conv1.0": "model.model.24.3.conv1.conv",
    "feature_map.backbone.layer4.4.conv1.0": "model.model.24.4.conv1.conv",
    "feature_map.backbone.layer4.5.conv1.0": "model.model.24.5.conv1.conv",
    "feature_map.backbone.layer4.6.conv1.0": "model.model.24.6.conv1.conv",
    "feature_map.backbone.layer4.7.conv1.0": "model.model.24.7.conv1.conv",
    "feature_map.backbone.layer4.0.conv2.0": "model.model.24.0.conv2.conv",
    "feature_map.backbone.layer4.1.conv2.0": "model.model.24.1.conv2.conv",
    "feature_map.backbone.layer4.2.conv2.0": "model.model.24.2.conv2.conv",
    "feature_map.backbone.layer4.3.conv2.0": "model.model.24.3.conv2.conv",
    "feature_map.backbone.layer4.4.conv2.0": "model.model.24.4.conv2.conv",
    "feature_map.backbone.layer4.5.conv2.0": "model.model.24.5.conv2.conv",
    "feature_map.backbone.layer4.6.conv2.0": "model.model.24.6.conv2.conv",
    "feature_map.backbone.layer4.7.conv2.0": "model.model.24.7.conv2.conv",
    "feature_map.backbone.conv20.0": "model.model.25.conv",
    "feature_map.backbone.conv21.0": "model.model.26.conv",
    "feature_map.backbone.conv22.0": "model.model.28.conv",
    "feature_map.backbone.conv23.0": "model.model.29.conv",
    "feature_map.backbone.conv24.0": "model.model.30.conv",
    "feature_map.backbone.layer5.0.conv1.0": "model.model.31.0.conv1.conv",
    "feature_map.backbone.layer5.1.conv1.0": "model.model.31.1.conv1.conv",
    "feature_map.backbone.layer5.2.conv1.0": "model.model.31.2.conv1.conv",
    "feature_map.backbone.layer5.3.conv1.0": "model.model.31.3.conv1.conv",
    "feature_map.backbone.layer5.0.conv2.0": "model.model.31.0.conv2.conv",
    "feature_map.backbone.layer5.1.conv2.0": "model.model.31.1.conv2.conv",
    "feature_map.backbone.layer5.2.conv2.0": "model.model.31.2.conv2.conv",
    "feature_map.backbone.layer5.3.conv2.0": "model.model.31.3.conv2.conv",
    "feature_map.backbone.conv25.0": "model.model.32.conv",
    "feature_map.backbone.conv26.0": "model.model.33.conv",
    "feature_map.backbone.conv27.0": "model.model.35.conv",
}


def convert_weight(ori_weight, new_weight):
    new_ckpt = []
    param_dict = ms.load_checkpoint(ori_weight)
    for k, v in param_dict.items():
        if "feature_map.backbone" in k:
            for key, val in convert_dict.items():
                if key in k:
                    k = k.replace(key, val)
        new_ckpt.append({"name": k, "data": v})
    ms.save_checkpoint(new_ckpt, new_weight)


if __name__ == "__main__":
    convert_weight(
        "./cspdarknet53_ascend_v120_imagenet2012_official_cv_bs64_top1acc7854_top5acc9428.ckpt",
        "./yolov4_backbone.ckpt",
    )
